import os
import shutil
import csv
import random
from collections import defaultdict

# 设置随机种子以确保可重复性
random.seed(42)

# 定义路径
BASE_DIR = 'samples/subtrain/clean'
CSV_PATH = 'samples/subtrainLabels.csv'

train_dir = os.path.join(BASE_DIR, 'train')
test_dir = os.path.join(BASE_DIR, 'test')

# 创建训练和测试目录
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)

# 读取CSV文件并建立文件名到类别的映射
file_class_map = {}
with open(CSV_PATH, 'r') as f:
    reader = csv.DictReader(f)
    for row in reader:
        file_class_map[row['Id'] + '.asm'] = int(row['Class'])

# 按类别分组文件
class_files = defaultdict(list)
for file in os.listdir(BASE_DIR):
    if file.endswith('.asm') and file in file_class_map:
        class_id = file_class_map[file]
        class_files[class_id].append(file)

# 对每个类别进行分层抽样
for class_id, files in class_files.items():
    random.shuffle(files)  # 随机打乱文件列表
    
    # 计算测试集数量（约20%）
    test_count = max(1, round(len(files) * 0.2))
    
    # 分割文件
    test_files = files[:test_count]
    train_files = files[test_count:]
    
    # 移动测试集文件
    for file in test_files:
        src = os.path.join(BASE_DIR, file)
        dst = os.path.join(test_dir, file)
        shutil.move(src, dst)
    
    # 移动训练集文件
    for file in train_files:
        src = os.path.join(BASE_DIR, file)
        dst = os.path.join(train_dir, file)
        shutil.move(src, dst)

print("数据集划分完成！")
print(f"训练集位置: {train_dir}")
print(f"测试集位置: {test_dir}")